from keras.optimizers import (
    Adam,
    Nadam,
)

from tensorflow import (
    Tensor,
    where,
    abs as tf_abs,
    subtract as tf_subtract,
    multiply as tf_multiply,
    square as tf_square,
    reduce_mean as tf_reduce_mean
)


def optimizer_adam(
        learning_rate: float = 1e-3,
        beta_1: float = 0.9,
        beta_2: float = 0.999,
        epsilon: float = 1e-7,
        amsgrad: bool = False
) -> Adam:
    """
    Adam优化器（2015）结合了AdaGrad（2010）和RMSProp（2012）的优点。

    Adam自适应地调整每个参数的学习率，基于参数梯度的
    一阶和二阶估计。它计算梯度和梯度平方的指数衰减
    移动平均，使用的衰减因子是beta_1（一阶）和beta_2（二阶）。

    更新公式：
        x_t = x_t-1 - lr * grad_mean_hat_t / sqrt(grad_var_hat_t)
        其中
        grad_mean_hat_t = beta_1 * grad_mean_hat_t-1 + (1 - beta_1) * grad_t
        grad_var_hat_t  = beta_2 * grad_var_hat_t-1 + (1 - beta_2) * grad_t^2

    AMSgrad是2018年的Adam变体，它通过保持最大方差估计来修正
    收敛行为中的误差：
        x_t = x_t-1 - lr * grad_mean_hat_t / sqrt(max_grad_var_hat_t)
        其中
        max_grad_var_hat_t = max(grad_var_hat_t, max_grad_var_hat_t-1)

    参数:
        learning_rate: SGD的步长
        beta_1: 一阶矩估计的衰减因子
        beta_2: 二阶矩估计的衰减因子
        epsilon: 用于数值稳定性的常量
        amsgrad: 是否使用AMSgrad变体
    """

    optimizer = Adam(
        learning_rate=learning_rate,  # 默认值1e-3
        beta_1=beta_1,  # 默认值0.9
        beta_2=beta_2,  # 默认值0.999
        epsilon=epsilon,  # 默认值1e-7，tf文档建议0.1或1，比较特殊
        amsgrad=amsgrad  # 默认值为False
    )

    return optimizer


def optimizer_nadam(
        learning_rate: float = 1e-3,
        beta_1: float = 0.9,
        beta_2: float = 0.999,
        epsilon: float = 1e-7
) -> Nadam:
    """
    NAdam（2016）实现了带Nesterov动量的Adam优化器。
    Nesterov动量在计算梯度之前应用动量估计，即：
    它在f(theta_t-1 - lr * beta_1 * grad_mean_estimate_t-1)上计算梯度。
    这通常会导致更快的收敛，并且具有一些证明的更好特性。

    参数:
        learning_rate: SGD的步长
        beta_1: 一阶矩估计的衰减因子
        beta_2: 二阶矩估计的衰减因子
        epsilon: 用于数值稳定性的常量
    """

    optimizer = Nadam(
        learning_rate=learning_rate,
        beta_1=beta_1,
        beta_2=beta_2,
        epsilon=epsilon
    )

    return optimizer


def mse_loss(
        td_error: Tensor,
        importance_sampling_weights: Tensor
) -> Tensor:
    """
    均方误差损失函数。

    计算带有重要性采样权重的均方误差（MSE），
    用于度量预测值和真实值之间的差距。

    参数:
        td_error: 时序差分误差
        importance_sampling_weights: 重要性采样权重

    返回:
        Tensor: 加权均方误差
    """
    squared_error = tf_multiply(importance_sampling_weights, tf_square(td_error))

    return tf_reduce_mean(squared_error)


def huber_loss(
        td_error: Tensor,
        importance_sampling_weights: Tensor
) -> Tensor:
    """
    Huber损失函数通过在损失曲线的线性部分修正梯度幅度来避免“梯度爆炸”。

    Huber损失尝试防止‘梯度爆炸’，并在误差较小时使用均方误差，
    误差较大时使用绝对误差。

    参数:
        td_error: 时序差分误差
        importance_sampling_weights: 重要性采样权重

    返回:
        Tensor: 加权Huber损失
    """
    delta = 1  # 在此处截断

    # 计算绝对误差
    absolute_error = tf_multiply(
        delta,
        tf_subtract(
            tf_abs(td_error),
            tf_multiply(.5, delta)
        )
    )

    # 计算平方误差
    squared_error = tf_multiply(
        .5,
        tf_square(td_error)
    )

    # 创建指示器，判断误差是否小于delta
    indicator = absolute_error < delta
    # 使用Huber损失公式：对于小误差使用平方误差，对于大误差使用绝对误差
    loss = where(indicator, squared_error, absolute_error)

    return tf_reduce_mean(importance_sampling_weights * loss)
